package com.zzyl.common.ai;

import com.baidubce.qianfan.Qianfan;
import com.baidubce.qianfan.core.auth.Auth;
import com.baidubce.qianfan.model.chat.ChatResponse;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.*;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.List;

@Component
public class AIModelInvoker {

    @Autowired
    private BaiduAIProperties baiDuProperties;

    public String chatWithQianfan(String prompt) {
        // 创建请求体
        QianfanRequest request = createRequest(prompt);

        // 发送请求并获取响应
        QianfanResponse response = sendRequest(request);

        // 提取并返回content内容
        return extractContent(response);
    }

    private QianfanRequest createRequest(String prompt) {
        // 创建消息列表
        List<Message> messages = new ArrayList<>();
        messages.add(new Message("system", "平台助手"));
        messages.add(new Message("user", prompt));

        // 创建请求对象
        return new QianfanRequest(baiDuProperties.getQianfanModel(), messages);
    }

    private QianfanResponse sendRequest(QianfanRequest request) {
        RestTemplate restTemplate = new RestTemplate();

        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", "Bearer " + baiDuProperties.getSecretKey());

        // 创建HTTP实体
        HttpEntity<QianfanRequest> entity = new HttpEntity<>(request, headers);

        try {
            // 发送POST请求
            ResponseEntity<QianfanResponse> response = restTemplate.exchange(
                    baiDuProperties.getAccessKey(),
                    HttpMethod.POST,
                    entity,
                    QianfanResponse.class
            );

            if (response.getStatusCode() == HttpStatus.OK) {
                return response.getBody();
            } else {
                System.err.println("请求失败，状态码: " + response.getStatusCode());
                return null;
            }
        } catch (Exception e) {
            System.err.println("请求异常: " + e.getMessage());
            e.printStackTrace();
            return null;
        }
    }

    private String extractContent(QianfanResponse response) {
        if (response == null || response.getChoices() == null || response.getChoices().isEmpty()) {
            return "抱歉，未能获取到有效响应";
        }

        // 获取第一个choice的message content
        Choice firstChoice = response.getChoices().get(0);
        if (firstChoice != null && firstChoice.getMessage() != null) {
            return firstChoice.getMessage().getContent();
        }

        return "抱歉，响应内容格式异常";
    }

    // 内部类 - 移除了static关键字
    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public  static class QianfanRequest {
        private String model;
        private List<Message> messages;
    }

    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public  static class Message {
        private String role;
        private String content;
    }

    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public  static class QianfanResponse {
        private String id;
        private String object;
        private Long created;
        private List<Choice> choices;
        private Usage usage;
    }

    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public  static class Choice {
        private Integer index;
        private Message message;
        private String finishReason;
    }

    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public  static class Usage {
        private Integer promptTokens;
        private Integer completionTokens;
        private Integer totalTokens;
    }
}
